#coding=utf-8


"""
Fitting the extremum of a light-curve of a variable, with a 2nd order polynominal.
Since 2022.10
By Dr Jie Zheng and Dr/Prof Lin-Qiao Jiang
"""


import numpy as np
from matplotlib import pyplot as plt
import re
import os
import sys
import datetime
import pltgui


#----------------------------------------------------------
# Transfer y=ax^2+bx+c to y = a(x-x0)^2+p
#----------------------------------------------------------
def parabola(co):
    a, b, c = co
    x0 = - b / a / 2
    p = c - b * b / a / 4
    return a, x0, p


#----------------------------------------------------------
# fit parabola Monte Carlo
#----------------------------------------------------------
def mcmc_fit_x2(bjd, bjderr, mag, magerr, ix, nn, mcfile=None):
    """
    Monte Carlo, and find the mean and stddev
    bjd, bjderr: barycentric jd and error
    mag, magerr: mag/flux and error
    ix: range of data
    nn: rounds of MC
    mcfile: mcmc result file
    returns: a, a_err, peak_bjd, peak_bjd_err, peak_mag, peak_mag_err
    """
    # array of results, cols: a、bjd_0、mag_0
    res = np.empty((nn, 3))
    
    for i in range(nn):
        # add normalized random errors
        bjdmc = bjd[ix] + np.random.normal(0, np.abs(bjderr[ix]))
        magmc = mag[ix] + np.random.normal(0, magerr[ix])
        # fit, transfer, and record
        co = np.polyfit(bjdmc, magmc, 2)
        res[i] = parabola(co)
    # log results
    if mcfile:
        np.savetxt(mcfile, res, fmt='%.6f')
    # get means and stddevs
    a = np.mean(res[:, 0])
    a_err = np.std(res[:, 0])
    peak_bjd = np.mean(res[:, 1])
    peak_bjd_err = np.std(res[:, 1])
    peak_mag = np.mean(res[:, 2])
    peak_mag_err = np.std(res[:, 2])

    return a, a_err, peak_bjd, peak_bjd_err, peak_mag, peak_mag_err


#----------------------------------------------------------
# Main Function
#----------------------------------------------------------
def i_para_fit(lc_file, mag_flux="Flux"):
    """
    lc_file: light-curve file
    mag_flux: Flux or Mag
    """
    
    #----------------------------------------------------------
    # Load data
    #----------------------------------------------------------
    bjd, bjderr, mag, magerr = np.loadtxt(
        lc_file,
        usecols=[4,5,2,3],
        unpack=True,
    )
    # error from tess is in seconds, convert to days
    bjderr /= 86400
    # Flux to mag
    if mag_flux.lower() == "mag":
        magerr = magerr / mag
        mag = -2.5 * np.log10(mag)
    
    # remove a base BJD, make data smaller
    bjd0 = int(min(bjd))
    # Sort by time, in case they are random
    ix = np.argsort(bjd)
    bjd = bjd[ix] - bjd0
    bjderr = bjderr[ix]
    mag = mag[ix]
    magerr = magerr[ix]
    
    # Interval between 2 item, used in plotting
    bjd_sp = np.median(bjd[1:] - bjd[:-1])
    
    #----------------------------------------------------------
    # Init result var
    #----------------------------------------------------------
    # fitting range, left and right
    tleft, tright  = None, None
    # parabola
    paxp = [None]*3
    resid_mean, resid_std = None, None
    
    #----------------------------------------------------------
    # GUI Init
    #----------------------------------------------------------
    # Axes, lignt-curve on top and buttons down
    fig = plt.figure()
    axlc = fig.add_axes([0.075,0.20, 0.88, 0.75])
    axbt = fig.add_axes([0.01, 0.01, 0.98, 0.09])
    
    # Plotting the whole lc
    z_lcall = axlc.plot(bjd, mag, "k.")
    # reverse y for mag
    if mag_flux == "mag":
        axlc.invert_yaxis()
    
    # All vars starts with 'z' means image parts
    # Fitting formular and residual
    z_formula = axbt.text(-12.5, -1.5, f'Flux = $a(bjd-bjd_0)^2+p$  $Resid\\times 10^3=\\phi \\pm \\sigma$', ha="center", va="center")
    
    # Left and right border
    z_lim_l = axlc.axvline([bjd[ 0]], color="m", linestyle="--")
    z_lim_r = axlc.axvline([bjd[-1]], color="m", linestyle="--")
    z_val_l = axlc.text(bjd[ 0], mag[ 0], 'X', ha="left", va="center")
    z_val_r = axlc.text(bjd[-1], mag[-1], 'X', ha="right", va="center")
    z_lim_l.set_visible(False)
    z_lim_r.set_visible(False)
    z_val_l.set_visible(False)
    z_val_r.set_visible(False)
    # points selected to fit, empty now
    z_lc_sel, = axlc.plot([], [], "g.")
    # fitting result
    z_lc_fit2, = axlc.plot([], [], "b:")
    # peak position
    z_peak, = axlc.plot([], [], 'r+')
    z_peak_bjd = axlc.text(bjd[0], mag[0], '', ha="center", va="top")
    z_peak_mag = axlc.text(bjd[0], mag[0], '', ha="left", va="center")
    
    axlc.set_title(f"{lc_file}")
    axlc.set_xlabel(f"BJD - {bjd0}")
    axlc.set_ylabel(mag_flux)
    
    #----------------------------------------------------------
    # Redraw left and right
    #----------------------------------------------------------
    def redraw_t():
        if tleft is not None:
            # Plot left border, find the nearest point
            i = np.argmin(np.abs(bjd - tleft))
            # set visible
            z_lim_l.set_visible(True)
            z_val_l.set_visible(True)
            # move the v-line and data marker
            z_lim_l.set_xdata([tleft])
            z_val_l.set_position((tleft, mag[i]))
            z_val_l.set_text(f"{tleft:.5f}")
        if tright is not None:
            # same as left
            i = np.argmin(np.abs(bjd - tright))
            z_lim_r.set_visible(True)
            z_val_r.set_visible(True)
            z_lim_r.set_xdata([tright]*2)
            z_val_r.set_position((tright, mag[i]))
            z_val_r.set_text(f"{tright:.5f}")
            
    #----------------------------------------------------------
    # Refit the curve, when left or right changed
    #----------------------------------------------------------
    def refit_x2():
        nonlocal paxp, resid_mean, resid_std
        # if not both left and right set, leave
        if tleft is None or tright is None:
            return
        # find points, if less than 5 points, leave
        ix = np.where((tleft <= bjd) & (bjd <= tright))[0]
        if len(ix) >= 5:
            # fit and convert
            co = np.polyfit(bjd[ix], mag[ix], 2)
            a, x0, p = parabola(co)
            paxp[:] = a, x0, p
            # residual
            resid = np.polyval(co, bjd[ix]) - mag[ix]
            resid_mean = np.mean(resid)
            resid_std = np.std(resid)
            # plot fitted curve
            bjd_fit = np.arange(tleft-bjd_sp, tright+bjd_sp*2, bjd_sp/2) # 适当外延，加大密度
            mag_fit = np.polyval(co, bjd_fit)
            z_lc_fit2.set_data(bjd_fit, mag_fit)
            # mark the points used
            z_lc_sel.set_data(bjd[ix], mag[ix])
            # the peak
            z_peak.set_data([x0], [p])
            z_peak_mag.set_position([x0+0.015, p])
            z_peak_bjd.set_position([x0, p-0.015])
            z_peak_mag.set_text(f"{p:.4f}")
            z_peak_bjd.set_text(f"{x0+bjd0:.5f}")
            # display result
            z_formula.set_text(f"$Flux = {a:.4f}(bjd-{x0+bjd0:.5f})^2{p:+.4f}$  "
                               f"$Resi\\times 10^3={resid_mean*1000:.1f}\\pm{resid_std*1000:.1f}$")

    #----------------------------------------------------------
    # Save parabola
    #----------------------------------------------------------
    def save_x2():
        # if the first time, create and add header
        if not os.path.isfile(f"{lc_file}.poly2.txt"):
            with open(f"{lc_file}.poly2.txt", "w") as ff:
                ff.write(
                    f"#    ProcDateTime        t_left  t_right  "
                    f"    a        bjd        mag      "
                    f"med_res   std_res   "
                    f"   mca       mcbjd     mcmag     "
                    f"  mcaer     mcbjder    mcmager \n"
                )
        
        # current time in ISO
        dt = datetime.datetime.now().strftime('%Y-%m-%dT%H-%M-%S')
        
        # Monte Carlo
        ix = np.where((tleft <= bjd) & (bjd <= tright))[0]
        a, a_err, peak_bjd, peak_bjd_err, peak_mag, peak_mag_err = \
            mcmc_fit_x2(bjd, bjderr, mag, magerr, ix, 1000, f"{lc_file}.{dt}.mcx2.txt")
        
        # Result string and output
        res = (
            f"{dt}   {tleft:.5f}   {tright:.5f}   "
            f"{paxp[0]:.3f}   {paxp[1]+bjd0:.5f}   {paxp[2]:+.4f}   "
            f"{resid_x2[0]:.3f}   {resid_x2[1]:.5f}   "
            f"{a:.3f}   {peak_bjd+bjd0:.5f}   {peak_mag:+.4f}   "
            f"{a_err:.3f}   {peak_bjd_err:.5f}   {peak_mag_err:+.4f}"
        )
        print(res)
        with open(f"{lc_file}.poly2.txt", "a") as ff:
            ff.write(res + "\n")
        # Save figure
        fig.savefig(f"{lc_file}.{dt}.png")
    

    def btnQ_click(btn):
        return True
    
    def btnL_click(btn=None):
        nonlocal tleft, tright
        # mark button left as checked
        btnL.check()
        # read a point from image axes
        p = fig.ginput()[0]
        if p[0] > 0:
            # if in, as new bound
            tleft = p[0]
            # if left-right flipped, swap
            if tright is not None and tleft > tright: tleft, tright = tright, tleft
            # redraw and refit
            redraw_t()
            refit_x2()
        btnL.uncheck()
    
    def btnR_click(btn=None):
        nonlocal tleft, tright
        # Same as left
        btnR.check()
        p = fig.ginput()[0]
        if p[0] > 0:
            tright = p[0]
            if tleft is not None and tleft > tright: tleft, tright = tright, tleft
            redraw_t()
            refit_x2()
        btnR.uncheck()
    
    def btnS_click(btn=None):
        # Save the result
        save_x2()
    
    # create button control
    btnctl = pltgui.i_btn_ctl(axbt)
    # colors for buttons
    btncl0 = '#cccccc'
    btncl1 = "y"
    btnclc = "g"
    # Use negative x for buttons
    btnL = btnctl.add_btn(-14,    0, 0.9, 2.0, btncl0, 'Left',    'L', btnL_click, chk_color=btncl1)
    btnR = btnctl.add_btn(-13,    0, 0.9, 2.0, btncl0, 'Right',   'R', btnR_click, chk_color=btncl1)
    btnS = btnctl.add_btn(-12,    0, 0.9, 2.0, btncl0, 'Save',    'S', btnS_click, chk_color=btncl1)
    btnQ = btnctl.add_btn(-11,    0, 0.9, 2.0, btncl0, 'Quit',    'Q', btnQ_click, chk_color=btncl1)
    # set x,y limit and turn off axis
    btnctl.set_axis_lim(0.025)
    axbt.set_ylim(-2, 1)

    plt.ion()
    btnctl.action_loop()
    plt.ioff()

    print("Bye")

if __name__ == "__main__":
    # Input parameters

    # At least a filename is required
    if len(sys.argv) < 2:
        print("""File name??""")
        exit()

    i_para_fit(sys.argv[1], "flux")
